//#include <pthread.h>
#include <mutex>
#include <thread>
//#include <unistd.h>
#include <signal.h>

#ifndef USE_LINUX
#include <winsock2.h>
#endif

#include <mysql.h>
#include <map>
#include "opencv2/opencv.hpp"

#include "cnn_function.hpp"
#include "util.h"

#include "rpc/server.h"

std::vector<std::string> current_user_name;
std::vector<int> current_user_spoofing;
std::vector<cv::Rect> rois_global;

// 0.62 for 50, 0.6 for 100
// sudo apt-get install libmysqlclient-dev
std::map<std::string, float**> all_group_feature;
std::map<std::string, int*> all_group_feature_id;
std::map<std::string, int> all_group_index;
int max_feature_num = 0;
std::map<int, int> all_delete_feature;
std::map<int, float> group_id_to_threshold;
float threshold_global = 0.6f;

bool have_init = false;
std::mutex gpu_lock;
bool recognition_stop = false;

int reload_feature();
int load_group_id_to_threshold();

void free_dnn_network(){
    uninit_detector();
    uninit_recognition();
    uninit_mtcnn();
}

void init_dnn_network(){
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    have_init = true;
    std::string cnn_dir = "/var/darknet/test/cnn";
    //init_recognition((cnn_dir + "/cfg/cosface_new_34.cfg").c_str(), (cnn_dir + "/model/face_30_glint_refine.cnn").c_str());
    init_recognition((cnn_dir + "/cfg/cosface_new_100.cfg").c_str(), (cnn_dir + "/model/model.cnn.100").c_str());
    init_detector((cnn_dir + "/cfg/yolov3-small.cfg").c_str(), (cnn_dir + "/model/yolov3-small_final_max_epoch_15.weights").c_str());
    init_mtcnn((cnn_dir + "/cfg/mtcnn_onet.cfg").c_str(), (cnn_dir + "/model/mtcnn_final.weights").c_str());
    //load_group_id_to_threshold();
    reload_feature();
}

int set_threshold(float threshold) {
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    threshold_global = threshold;
    printf("set_threshold %f\n", threshold);
    return 1;
}

void print_error_message(MYSQL *connection){
    if (mysql_errno(connection)) {
        fprintf(stderr, "mysql error %d: %s\n", mysql_errno(connection), mysql_error(connection));
    }
}

int load_group_id_to_threshold(){
    group_id_to_threshold.clear();
    MYSQL *connection = mysql_init(NULL);
    if (connection == NULL) {
        printf("mysql_init failed!\n");
        return -1;
    }
    connection = mysql_real_connect(connection, "localhost", "iim", "mysql@120test", "recognition", 0, NULL, 0);
    if (connection) {
        printf("Connection success! ");
    } else {
        print_error_message(connection);
        mysql_close(connection);
        return 1;
    }
    char query_sql[256] = {0};
    sprintf(query_sql, "select * from `group` where deleted = 0 order by id");
    int query_result = mysql_query(connection, query_sql);
    if(query_result){
        print_error_message(connection);
        mysql_close(connection);
        return 2;
    } else {
        MYSQL_RES *res_ptr = mysql_store_result(connection);
        if(res_ptr) {
            printf("query %lu Rows", (unsigned long)mysql_num_rows(res_ptr));
            unsigned int num_fields = mysql_num_fields(res_ptr);
            printf("num_fields: %u\n", num_fields);
            MYSQL_ROW row;
            while((row = mysql_fetch_row(res_ptr)))  {
                int group_id = -1;
                float threshold = -1;
                for(unsigned int i = 0; i < num_fields; i++){
                    if(i == 0){
                        group_id = atoi(row[i]);
                    } else if (3 == i){
                        threshold = atof(row[i]);
                    } else {
                    }
                }
                group_id_to_threshold[group_id] = threshold;
            }
            if (mysql_errno(connection)) {
                fprintf(stderr, "Retrive error:%s\n", mysql_error(connection));
            }
            mysql_free_result(res_ptr);
        } else {
            print_error_message(connection);
            mysql_close(connection);
            return 3;
        }
    }
    mysql_close(connection);
    std::map<int, float>::iterator it;
    for(it = group_id_to_threshold.begin(); it != group_id_to_threshold.end(); it++){
        printf("group_id_to_threshold: %d %f\n", it->first, it->second);
    }
    return 0;
}

int delete_feature(int feature_id){
    all_delete_feature[feature_id] = 1;
    printf("delete_feature: %d\n", feature_id);
    return 0;
}

int add_feature(int feature_id, std::string feature_str, std::string group_id){
    //printf("%d, %s %s", feature_id, group_id.c_str(), feature_str.c_str());
    if(all_group_index[group_id] < max_feature_num && all_group_feature.size() > 0){
        const char *float_str = feature_str.c_str();
        float *feature_float = (float *)malloc(512 * sizeof(float));
        for(unsigned int j = 0; j < 512; j++) {
            feature_float[j] = atof(float_str);
            float_str = strchr(float_str, ',') + 1;
        }

        all_group_feature_id[group_id][all_group_index[group_id]] = feature_id;
        all_group_feature[group_id][all_group_index[group_id]] = feature_float;
        all_group_index[group_id] += 1;
    } else {
        printf("add_feature: feature exceed max_feature_num reload_feature\n");
        reload_feature();
    }
    return 0;
}

int reload_feature() {
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    for(std::map<std::string, int*>::iterator it = all_group_feature_id.begin(); it != all_group_feature_id.end(); it++){
        free(it->second);
    }
    for(std::map<std::string, float**>::iterator itt = all_group_feature.begin(); itt != all_group_feature.end(); itt++){
        for(int i = 0; i < all_group_index[itt->first]; i++){
            free(itt->second[i]);
            //printf("%d: %d %f %f %f %f\n", i, itt->first, itt->second[i][0], itt->second[i][1], itt->second[i][510], itt->second[i][511]);
        }
        free(itt->second);
    }
    max_feature_num = 0;
    all_group_index.clear();
    all_delete_feature.clear();
    all_group_feature.clear();
    all_group_feature_id.clear();
    MYSQL *connection = mysql_init(NULL);
    if (connection == NULL) {
        printf("mysql_init failed!\n");
        return -1;
    }
    connection = mysql_real_connect(connection, "localhost", "iim", "mysql@120test", "recognition", 0, NULL, 0);
    if (connection) {
        printf("Connection success! ");
    } else {
        print_error_message(connection);
        mysql_close(connection);
        return 1;
    }
    char query_sql[256] = {0};
    //sprintf(query_sql, "select id, group_id, feature from person where status = 1 order by id"); // for android face34 model
    sprintf(query_sql, "select id, group_id, feature from feature where status = 1 order by id");
    int query_result = mysql_query(connection, query_sql);
    if(query_result){
        print_error_message(connection);
        mysql_close(connection);
        return 2;
    } else {
        MYSQL_RES *res_ptr = mysql_store_result(connection);
        if(res_ptr) {
            unsigned long row_num = mysql_num_rows(res_ptr);
            if(row_num < 500){
                max_feature_num = 1000;
            } else {
                max_feature_num = row_num * 2;
            }
            printf("query %lu Rows max_feature_num %d", row_num, max_feature_num);
            unsigned int num_fields = mysql_num_fields(res_ptr);
            printf("num_fields: %u\n", num_fields);
            MYSQL_ROW row;
            while((row = mysql_fetch_row(res_ptr)))  {
                //unsigned long *field_lengths = mysql_fetch_lengths(res_ptr);
                int feature_id = -1;
                std::string group_id;
                float *feature_float = (float *)malloc(512 * sizeof(float));
                for(unsigned int i = 0; i < num_fields; i++){
                    if(i == 0){
                        feature_id = atoi(row[i]);
                    } else if (1 == i){
                        group_id = row[i];
                    }else if(2 == i){
                        char *float_str = (char *)(row[i]);
                        for(unsigned int j = 0; j < 512; j++) {
                            feature_float[j] = atof(float_str);
                            float_str = strchr(float_str, ',') + 1;
                        }
                    } else {
                        printf("reload_feature: database field error\n");
                        exit(-1);
                    }
                }
                std::map<std::string, int*>::iterator it = all_group_feature_id.find(group_id);
                if (it == all_group_feature_id.end()){
                    float **all_group_feature_float_ptr = (float **)calloc(max_feature_num, sizeof(float *));
                    all_group_feature[group_id] = all_group_feature_float_ptr;
                    int *all_group_feature_id_ptr = (int *)calloc(max_feature_num, sizeof(int));
                    all_group_feature_id[group_id] = all_group_feature_id_ptr;
                    all_group_index[group_id] = 0;
                }
                all_group_feature_id[group_id][all_group_index[group_id]] = feature_id;
                all_group_feature[group_id][all_group_index[group_id]] = feature_float;
                all_group_index[group_id] += 1;
            }
            if (mysql_errno(connection)) {
                fprintf(stderr, "Retrive error:%s\n", mysql_error(connection));
            }
            mysql_free_result(res_ptr);
        } else {
            print_error_message(connection);
            mysql_close(connection);
            return 3;
        }
    }
    mysql_close(connection);
/*
    for(std::map<std::string, int*>::iterator it = all_group_feature_id.begin(); it != all_group_feature_id.end(); it++){
        for(int i = 0; i < all_group_index[it->first]; i++){
            std::cout << i << " " << it->first << " " << it->second[i] << std::endl;
        }
    }
*/
    for(std::map<std::string, int>::iterator it = all_group_index.begin(); it != all_group_index.end(); it++){
        std::cout << "all_group_index: " << it->first << " " << it->second << std::endl;
    }
/*
    for(std::map<std::string, float**>::iterator it = all_group_feature.begin(); it != all_group_feature.end(); it++){
        for(int i = 0; i < all_group_index[it->first]; i++){
            std::cout << "feature: " << i << " " <<  it->first << " " << it->second[i][0] << " " << it->second[i][511] << std::endl;
        }
    }
*/
    printf("all_group_feature size %lu\n", all_group_feature.size());
    printf("reload_feature over\n");
    return 0;
}

std::string registration_image(std::vector<char> const& img) {
    double start = getMsecOfNow();
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);

    int code = 0;
    if(!have_init){
        init_dnn_network();
    }
    int detection_bbox[MAX_BBOX_NUM * 4];

    std::vector<cv::Mat> aligned_faces;
    int face_features_length = 512;
    int face_count = 0;
    detection_face(image_input, &face_count, detection_bbox);

    /*
    for(int i = 0; i < face_count; ++i){
        cv::Rect roi_local(cv::Point(detection_bbox[i * 4], detection_bbox[i * 4 + 1]),
                           cv::Size(detection_bbox[i * 4 + 2] - detection_bbox[i * 4],
                                    detection_bbox[i * 4 + 3] - detection_bbox[i * 4 + 1]));
        printf("face size %d %d\n", roi_local.width, roi_local.height);
        cv::rectangle(image_input, cvPoint(roi_local.x, roi_local.y),
                      cvPoint(roi_local.x + roi_local.width, roi_local.y + roi_local.height), cvScalar(255,0,0), 4);
    }
    static int save_index = 0;
    cv::imwrite("/home/iim/testss" + std::to_string(save_index) + ".jpg", image_input);
    save_index += 1;
    */

    if(face_count == 0){
        code = 1;  // no face
        return std::to_string(code);
    } else if(face_count > 1){
        //code = 2;  // multi face
        //return std::to_string(code);
        int first_face_size = (detection_bbox[2] - detection_bbox[0]) * (detection_bbox[3] - detection_bbox[1]);
        for(int i = 1; i < face_count; ++i){
            int face_size = (detection_bbox[i * 4 + 2] - detection_bbox[i * 4]) * (detection_bbox[i * 4 + 3] - detection_bbox[i * 4 + 1]);
            if(first_face_size < face_size){
                first_face_size = face_size;
                for(int j = 0; j < 4; ++j) detection_bbox[j] = detection_bbox[i * 4 + j];
            }
        }
        std::cout << "registration_image face_count " << face_count << ", use max size face" << std::endl;
        face_count = 1;
    }
    if(detection_bbox[2] - detection_bbox[0] < 48 || detection_bbox[3] - detection_bbox[1] < 56){
        code = 3;  // face size small
        return std::to_string(code);
    }
    float *face_feature = (float *)malloc(face_count * face_features_length *sizeof(float));
    extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 1);
    std::string registration_image_feature_str;
    if(code == 0){
        std::string feature_str = "";
        for(int i = 0; i < 512; i++){
            feature_str += std::to_string(face_feature[i]) + ",";
        }
        registration_image_feature_str = feature_str;
    } else {
        registration_image_feature_str = std::to_string(code);
    }
    free(face_feature);
    std::cout << "registration_image spend: " << getMsecOfNow() - start << " ms code " << code << ", face_count " << face_count
              << " image width " << image_input.cols << " image width " << image_input.rows
              <<  " thread_id " <<  std::this_thread::get_id() << std::endl;
    return registration_image_feature_str;
}

std::string recognition_face(std::vector<char> const& img, std::string group_id) {
    if(!have_init){
        init_dnn_network();
    }
    int code = 0;
    int face_count = 0;
    std::string bbox_str = "";
    std::string feature_ids = "";
    int detection_bbox[MAX_BBOX_NUM * 4];
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    double start = getMsecOfNow();
    int face_features_length = 512;
    float *face_feature;
    {
        std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
        detection_face(image_input, &face_count, detection_bbox);
        //for(int i = 0; i < 4; i++) printf("%d %d\n", i, detection_bbox[i]);
        if(face_count == 0){
            code = 1;  // no face
            //free(face_feature);
            return std::to_string(code) + "|" + std::to_string(face_count) + "|" + feature_ids + "|" + bbox_str + "|" + std::to_string(threshold_global);
        }
        face_feature = (float *)malloc(face_count * face_features_length * sizeof(float));
        extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 0);
    }

    std::vector<float **> compare_group_feature;
    std::vector<int> compare_group_feature_num;
    std::vector<std::string> compare_group_group_id;
    std::map<std::string, int*>::iterator all_group_feature_id_it = all_group_feature_id.find(group_id);
    bool compare_all_feature = false;
    if (group_id == "" || all_group_feature_id_it == all_group_feature_id.end()){
        compare_all_feature = true;
        printf("not have group_id %s\n", group_id.c_str());
        for(std::map<std::string, float**>::iterator it = all_group_feature.begin(); it != all_group_feature.end(); it++){
            compare_group_feature.push_back(it->second);
            compare_group_feature_num.push_back(all_group_index[it->first]);
            compare_group_group_id.push_back(it->first);
        }
    } else {
        printf("have group_id %s\n", group_id.c_str());
        float **group_feature = all_group_feature[group_id];
        int group_feature_num = all_group_index[group_id];
        compare_group_feature.push_back(group_feature);
        compare_group_feature_num.push_back(group_feature_num);
    }

    for(int j = 0; j < face_count; j++){
        float max_score = -1.0F;
        int max_score_feature_id = -1;
        std::string max_score_group_id;
        float second_score = -1.0F;
        int second_score_feature_id = -1;
        std::string second_score_group_id;
        float third_score = -1.0F;
        int third_score_feature_id = -1;
        std::string third_score_group_id;
        for(unsigned int n = 0; n < compare_group_feature.size(); n++){
            for(int i = 0; i < compare_group_feature_num[n]; i++){
                if(all_delete_feature.size() > 0){
                    if(compare_all_feature) {
                        std::map<int, int>::iterator it_local = all_delete_feature.find(all_group_feature_id[compare_group_group_id[n]][i]);
                        if(it_local != all_delete_feature.end()) {
                            //printf("continue feature %d %d", it_local->first, it_local->second);
                            continue;
                        }
                    } else {
                        std::map<int, int>::iterator it_local = all_delete_feature.find(all_group_feature_id[group_id][i]);
                        if(it_local != all_delete_feature.end()) {
                            //printf("continue feature %d %d", it_local->first, it_local->second);
                            continue;
                        }
                    }
                }
                float score = 0;
                float *register_feature = compare_group_feature[n][i];
                for(int m = 0; m < face_features_length; m++){
                    score += face_feature[j * face_features_length + m] * register_feature[m];
                }
                if(score > max_score){
                    third_score = second_score;
                    third_score_feature_id = second_score_feature_id;
                    second_score = max_score;
                    second_score_feature_id = max_score_feature_id;
                    max_score = score;
                    max_score_feature_id = i;
                    if(compare_all_feature) {
                        third_score_group_id = second_score_group_id;
                        second_score_group_id = max_score_group_id;
                        max_score_group_id = compare_group_group_id[n];
                    }
                } else if(score > second_score){
                    third_score = second_score;
                    third_score_feature_id = second_score_feature_id;
                    second_score = score;
                    second_score_feature_id = i;
                    if(compare_all_feature) {
                        third_score_group_id = second_score_group_id;
                        second_score_group_id = compare_group_group_id[n];
                    }
                } else if(score > third_score){
                    third_score = score;
                    third_score_feature_id = i;
                    if(compare_all_feature) third_score_group_id = compare_group_group_id[n];
                }
            }
        }
        if(compare_all_feature){
            feature_ids += (max_score_feature_id != -1 ? std::to_string(all_group_feature_id[max_score_group_id][max_score_feature_id]) : "-1") +
                "," + std::to_string(max_score) + "," +
                (second_score_feature_id != -1 ? std::to_string(all_group_feature_id[second_score_group_id][second_score_feature_id]) : "-1") +
                "," + std::to_string(second_score) + "," +
                (third_score_feature_id != -1 ? std::to_string(all_group_feature_id[third_score_group_id][third_score_feature_id]) : "-1") +
                "," + std::to_string(third_score) + ",";
        } else {
            feature_ids += (max_score_feature_id != -1 ? std::to_string(all_group_feature_id[group_id][max_score_feature_id]) : "-1") +
                "," + std::to_string(max_score) + "," +
                (second_score_feature_id != -1 ? std::to_string(all_group_feature_id[group_id][second_score_feature_id]) : "-1") +
                "," + std::to_string(second_score) + "," +
                (third_score_feature_id != -1 ? std::to_string(all_group_feature_id[group_id][third_score_feature_id]) : "-1") +
                "," + std::to_string(third_score) + ",";
        }
        for(int i = 0; i < 4; i++){
            bbox_str += std::to_string(detection_bbox[j * 4 + i]) + ",";
        }
    }
    std::cout << "recognition_face spend: " << getMsecOfNow() - start << " ms code, face_count " << face_count
              <<  " thread_id " <<  std::this_thread::get_id() << std::endl;
    free(face_feature);
    return std::to_string(code) + "|" + std::to_string(face_count) + "|" + feature_ids + "|" + bbox_str + "|" +  std::to_string(threshold_global);
}

void sig_close(int sig){
    const char *s = NULL;
    switch (sig) {
    case SIGINT:
        s = "interrupt";
        break;
    case SIGTERM:
        s = "terminate";
        break;
    }
    recognition_stop = true;
    printf("\n\nreceived signal: %s, quit...\n", s);
}

int main(int argc, char **argv)
{
    signal(SIGINT, sig_close);  //close CTRL+C
    signal(SIGTERM, sig_close);

    rpc::server srv("127.0.0.1", 18080);
    srv.bind("registration_image", &registration_image);
    srv.bind("recognition_face", &recognition_face);
    srv.bind("set_threshold", set_threshold);
    srv.bind("reload_feature", reload_feature);
    srv.bind("add_feature", add_feature);
    srv.bind("delete_feature", delete_feature);
    //srv.run();
    srv.async_run();
    while (!recognition_stop){
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    printf("rpc quit...\n");
    if(have_init){
        free_dnn_network();
    }
    return 0;
}
